import copy
import numpy as np
import matplotlib.pyplot as plt

from path import PathButterfly
from double_s import DoubleS
from trajectory import Trajectory
from robot import Robot


class DynamicsScaling:
    def __init__(self, robot: Robot, trajectory: Trajectory, torque_limit: np.ndarray, num: int = 1001):
        self.trajectory = copy.deepcopy(trajectory)
        self.robot = copy.deepcopy(robot)
        self.num = num

        static_torque = np.zeros((self.num, 2))
        dynamic_torque = np.zeros((self.num, 2))

        i = 0
        for ti in np.linspace(0, self.trajectory.get_t(), self.num):
            trajectory = self.trajectory.get_trajectory(ti)

            robot_static_torque = robot.get_static_torque(trajectory[0], False)
            robot_dynamic_torque = robot.get_dynamic_torque(*trajectory, False)
            static_torque[i, :] = robot_static_torque
            dynamic_torque[i, :] = robot_dynamic_torque
            i = i + 1

        torque_max_no_static = torque_limit - static_torque
        torque_min_no_static = -torque_limit - static_torque

        self.factor = 1.0 / np.sqrt(np.max((np.max(dynamic_torque / torque_max_no_static), np.max(dynamic_torque / torque_min_no_static))))

    def get_trajectory(self, t):
        trajectory = self.trajectory.get_trajectory(t * self.factor)
        trajectory_double_s = self.trajectory.get_trajectory_double_s(t * self.factor)
        torque = self.robot.get_torque(trajectory[0], [i * self.factor for i in trajectory[1]],
                                       [i * self.factor ** 2 for i in trajectory[2]], False)
        return (trajectory_double_s[0],
                trajectory_double_s[1] * self.factor,
                trajectory_double_s[2] * self.factor ** 2,
                trajectory_double_s[3] * self.factor ** 3,
                torque,
                trajectory[0]
                )

    def get_total_time(self):
        return self.trajectory.get_t() / self.factor


if __name__ == '__main__':
    q0 = np.array(0)
    q1 = np.array(np.pi * 2)

    vmax = np.array(2.0)
    amax = np.array(5.0)
    jmax = np.array(20.0)
    torque_limit = np.array([100.0, 100])

    num = 1001

    path_butterfly = PathButterfly(0.2)
    double_s = DoubleS(q0, q1, vmax, amax, jmax)
    trajectory = Trajectory(path_butterfly, double_s)

    robot = Robot()
    dynamics_scaling = DynamicsScaling(robot, trajectory, torque_limit, num)
    print('double S total time: ', double_s.get_total_time())
    print('dynamics scaling total time: ', dynamics_scaling.get_total_time())

    q_double_s = []
    dq_double_s = []
    ddq_double_s = []
    dddq_double_s = []
    torque_double_s = []

    q_dynamics_scaling = []
    dq_dynamics_scaling = []
    ddq_dynamics_scaling = []
    dddq_dynamics_scaling = []
    torque_dynamics_scaling = []

    qi = []

    t = np.linspace(0, np.ceil(np.max((double_s.get_total_time(), dynamics_scaling.get_total_time()))), num)

    for ti in t:
        trajectory_double_s = double_s.get_trajectory(ti)
        q_double_s.append(trajectory_double_s[0])
        dq_double_s.append(trajectory_double_s[1])
        ddq_double_s.append(trajectory_double_s[2])
        dddq_double_s.append(trajectory_double_s[3])
        torque_double_s.append(robot.get_torque(*trajectory.get_trajectory(ti), False))

        trajectory_dynamics_scaling = dynamics_scaling.get_trajectory(ti)
        q_dynamics_scaling.append(trajectory_dynamics_scaling[0])
        dq_dynamics_scaling.append(trajectory_dynamics_scaling[1])
        ddq_dynamics_scaling.append(trajectory_dynamics_scaling[2])
        dddq_dynamics_scaling.append(trajectory_dynamics_scaling[3])
        torque_dynamics_scaling.append(trajectory_dynamics_scaling[4])

        qi.append(robot.ikine(trajectory_dynamics_scaling[5]))


    plt.figure(1)
    plt.subplot(2, 2, 1)
    plt.plot(t, q_double_s, label='double S')
    plt.plot(t, q_dynamics_scaling, label='dynamics scaling')
    plt.xlabel(r'$t (s)$')
    plt.ylabel(r'$q (rad)$')
    plt.legend()

    plt.subplot(2, 2, 2)
    plt.plot(t, dq_double_s, label='double S')
    plt.plot(t, dq_dynamics_scaling, label='dynamics scaling')
    plt.xlabel(r'$t (s)$')
    plt.ylabel(r'$\dot{q} (rad/s)$')
    plt.legend()

    plt.subplot(2, 2, 3)
    plt.plot(t, ddq_double_s, label='double S')
    plt.plot(t, ddq_dynamics_scaling, label='dynamics scaling')
    plt.xlabel(r'$t (s)$')
    plt.ylabel(r'$\ddot{q} (rad/s^2)$')
    plt.legend()

    plt.subplot(2, 2, 4)
    plt.plot(t, dddq_double_s, label='double S')
    plt.plot(t, dddq_dynamics_scaling, label='dynamics scaling')
    plt.xlabel(r'$t (s)$')
    plt.ylabel(r'$\dddot{q} (rad/s^3)$')
    plt.legend()

    plt.tight_layout()

    plt.figure(2)
    plt.subplot(2, 1, 1)
    plt.plot(t, [*zip(*torque_double_s)][0], label='double S')
    plt.plot(t, [*zip(*torque_dynamics_scaling)][0], label='dynamics scaling')
    plt.axhline(torque_limit[0], label='max torque')
    plt.axhline(-torque_limit[0], label='min torque')
    plt.xlabel(r'$t (s)$')
    plt.ylabel(r'$\tau (Nm)$')
    plt.legend()

    plt.subplot(2, 1, 2)
    plt.plot(t, [*zip(*torque_double_s)][1], label='double S')
    plt.plot(t, [*zip(*torque_dynamics_scaling)][1], label='dynamics scaling')
    plt.axhline(torque_limit[1], label='max torque')
    plt.axhline(-torque_limit[1], label='min torque')
    plt.xlabel(r'$t (s)$')
    plt.ylabel(r'$\tau (Nm)$')
    plt.legend()

    plt.tight_layout()
    plt.show()

    robot.plot(np.array(qi))
